#include "execution/executors/topn_executor.h"
#include <utility>

namespace bustub {

TopNExecutor::TopNExecutor(ExecutorContext *exec_ctx, const TopNPlanNode *plan,
                           std::unique_ptr<AbstractExecutor> &&child_executor)
    : AbstractExecutor(exec_ctx),
      plan_(plan),
      child_executor_(std::move(child_executor)),
      top_n_heap_(TupleCompare(plan_)) {}

void TopNExecutor::Init() {
  child_executor_->Init();
  TupleCompare comparator(plan_);
  Tuple tuple{};
  RID dummy_rid{};
  while (top_n_heap_.size() < plan_->GetN()) {
    if (child_executor_->Next(&tuple, &dummy_rid)) {
      top_n_heap_.push(std::move(tuple));
    } else {
      while (!top_n_heap_.empty()) {
        auto t = top_n_heap_.top();
        top_n_heap_.pop();
        res_.push_back(std::move(t));
      }
      return;
    }
  }
  while (child_executor_->Next(&tuple, &dummy_rid)) {
    if (comparator(tuple, top_n_heap_.top())) {
      top_n_heap_.pop();
      top_n_heap_.push(tuple);
    }
  }
  while (!top_n_heap_.empty()) {
    auto t = top_n_heap_.top();
    top_n_heap_.pop();
    res_.push_back(std::move(t));
  }
}

auto TopNExecutor::Next(Tuple *tuple, RID *rid) -> bool {
  if (res_.empty()) {
    return false;
  }
  *tuple = std::move(res_.back());
  res_.pop_back();
  return true;
}

auto TopNExecutor::GetNumInHeap() -> size_t { return top_n_heap_.size(); };

}  // namespace bustub
